{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mp.weixin.qq.com/s?__biz=MzkyMDMxNDkxOA==&mid=2247484467&idx=1&sn=f6ed20b751585bc0f355b2d60c64517d&chksm=c195f17af6e2786cb6d55c8814aba8967d76f41c13e47af5b5dc19da95d61b4b0e536f4cae72&scene=178&cur_album_id=3447886516519747585#rd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 振动升降指标ASI策略.py\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hist_k_data(code,start_date,end_date,frequency='d')->pd.DataFrame:\n",
    "    \"\"\"\n",
    "    获取历史K线数据\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    import baostock as bs\n",
    "    bs.login()\n",
    "    rs = bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,pctChg\",start_date,end_date,frequency=frequency)\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    result.open = result.open.astype(float)\n",
    "    result.high = result.high.astype(float)\n",
    "    result.low = result.low.astype(float)\n",
    "    result.close = result.close.astype(float)\n",
    "    result.date= pd.to_datetime(result.date)\n",
    "    result.set_index('date',inplace=True)\n",
    "    bs.logout()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取股票数据\n",
    "stock_symbol = 'sz.000001'\n",
    "start_date = '2005-01-01'\n",
    "end_date = '2023-01-01'\n",
    "data = get_hist_k_data(stock_symbol, start_date, end_date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 计算SAR指标\n",
    "def calculate_sar(data, af_start=0.02, af_max=0.20):\n",
    "    data['SAR'] = np.nan\n",
    "    data['AF'] = af_start\n",
    "    data['EP'] = data['high'].rolling(2).max()\n",
    "    data['SAR'][0] = data['high'][0]\n",
    "    for i in range(1, len(data)):\n",
    "        if data['close'][i] > data['SAR'][i-1]:\n",
    "            data['SAR'][i] = data['SAR'][i-1] + data['AF'][i-1] * (data['EP'][i-1] - data['SAR'][i-1])\n",
    "            data['EP'][i] = data['high'][i]\n",
    "            data['AF'][i] = min(data['AF'][i-1] + af_start, af_max)\n",
    "        elif data['close'][i] < data['SAR'][i-1]:\n",
    "            data['SAR'][i] = data['SAR'][i-1] - data['AF'][i-1] * (data['SAR'][i-1] - data['EP'][i-1])\n",
    "            data['EP'][i] = data['low'][i]\n",
    "            data['AF'][i] = min(data['AF'][i-1] + af_start, af_max)\n",
    "        else:\n",
    "            data['SAR'][i] = data['SAR'][i-1]\n",
    "            data['AF'][i] = data['AF'][i-1]\n",
    "    return data\n",
    "\n",
    "# 计算SAR指标\n",
    "data = calculate_sar(data)\n",
    "\n",
    "# 生成交易信号\n",
    "data['Signal'] = 0\n",
    "data['Position'] = 0\n",
    "data['Position'] = data['close'].shift(1) > data['SAR']\n",
    "data['Position'] = data['Position'].shift(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 回测策略\n",
    "def backtest_strategy(data):\n",
    "    data['Strategy_Returns'] = data['Position'].shift(1) * data['close'].pct_change()\n",
    "    data['Cumulative_Returns'] = (1 + data['Strategy_Returns']).cumprod()\n",
    "    return data\n",
    "\n",
    "data = backtest_strategy(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 绘制策略的累计收益\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(data['Cumulative_Returns'], label='SAR Strategy')\n",
    "plt.title('SARAO Strategy Cumulative Returns')\n",
    "plt.xlabel('Date')\n",
    "plt.ylabel('Cumulative Returns')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
